
import os
import subprocess
import glob

# Loop through folders 1-28
for folder_num in range(1, 28):
    folder_name = str(folder_num)
    
    if os.path.isdir(folder_name):
        print(f"Processing folder {folder_num}")
        
        # Change to the folder
        os.chdir(folder_name)
        
        # Find all *_p_aligned_tf##.nii files
        input_files = glob.glob("*_p_aligned_tf*.nii")
        
        if not input_files:
            print(f"  No input files found in folder {folder_num}")
            os.chdir('..')
            continue
        
        # Create atlas folder if it doesn't exist
        atlas_dir = 'atlas'
        if not os.path.exists(atlas_dir):
            os.makedirs(atlas_dir)
            print(f"  Created atlas directory: {atlas_dir}")
        
        # Process each input file
        for input_filename in input_files:
            # Extract subject ID (C001, C002, etc.)
            subject = input_filename.split('_')[0]
            
            # Extract sample number from the input filename
            # Assuming format: C001_sample3_p_aligned_tf##.nii
            sample_part = input_filename.split('_')[1]  # Gets "sample3" or "sample4"
            
            # Create output filename in the atlas folder
            output_basename = input_filename.replace('_p_aligned_tf', '_p_registered_tf')
            output_filename = os.path.join(atlas_dir, output_basename)
            
            # Define transformation file paths (in parent directory)
            # Use the same sample number as in the input file
            template_path = '../Hamper_tf1_template.nii.gz'
            warp_path = f'../Hamper_tf1_{subject}_{sample_part}_p_aligned_tf01Warp.nii.gz'
            affine_path = f'../Hamper_tf1_{subject}_{sample_part}_p_aligned_tf01Affine.txt'
            
            print(f"  Processing {subject}: {input_filename} -> {output_filename}")
            
            # Check if transformation files exist
            if not os.path.isfile(template_path):
                print(f"    ✗ Template file missing: {template_path}")
                continue
            if not os.path.isfile(warp_path):
                print(f"    ✗ Warp file missing: {warp_path}")
                continue
            if not os.path.isfile(affine_path):
                print(f"    ✗ Affine file missing: {affine_path}")
                continue
            
            # Check WarpImageMultiTransform executable
            if not os.path.isfile('../WarpImageMultiTransform'):
                print("    ✗ WarpImageMultiTransform executable not found")
                break
            
            # Build command
            command = [
                '../WarpImageMultiTransform',
                '3',
                input_filename,
                output_filename,
                '-R',
                template_path,
                warp_path,
                affine_path
            ]
            
            print(f"    Command: {' '.join(command)}")
            
            try:
                # Run the command
                result = subprocess.run(command, check=True, capture_output=True, text=True)
                print(f"    ✓ Success!")
                
            except subprocess.CalledProcessError as error:
                print(f"    ✗ Failed with return code: {error.returncode}")
                if error.stderr:
                    print(f"      Error: {error.stderr}")
                if error.stdout:
                    print(f"      Output: {error.stdout}")
                    
            except FileNotFoundError:
                print("    ✗ WarpImageMultiTransform executable not found")
                break
        
        # Return to parent directory
        os.chdir('..')
        print()

print("Processing complete!")